import aiomysql
import asyncio
import json
from datetime import datetime
from collections import OrderedDict
import logging


class LawDB:
    def __init__(self):
        self.pool = None


    async def init_pool(self, retry_times=8, retry_delay=5):
        for attempt in range(retry_times):
            try:
                self.pool = await aiomysql.create_pool(
                    host='192.168.2.244',
                    port=3306,
                    user='root',
                    password='linkAge@12345',
                    db='stock_exchange',
                    charset='utf8mb4',
                    loop=asyncio.get_event_loop(),
                    minsize=1,
                    maxsize=800,
                    connect_timeout=300
                )
                logging.info("数据库连接池初始化成功")
                return
            except aiomysql.MySQLError as e:
                logging.error(f"数据库连接池初始化失败，尝试第 {attempt + 1} 次重试: {e}")
                await asyncio.sleep(retry_delay)
        logging.error("数据库连接池初始化失败，重试次数已达上限")
        raise ConnectionError("无法连接到数据库")

    async def close_pool(self):
        if self.pool:
            self.pool.close()
            await self.pool.wait_closed()
            logging.info("数据库连接池已关闭")

    # 查询 word2json表
    async def select_word2json(self, fileId: str):

        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                select_sql = """
                           SELECT word_json FROM word2json WHERE fileId = %s
                           """
                try:
                    await cursor.execute(select_sql, (fileId))
                    json_results = await cursor.fetchone()
                    print("word2json表 查询成功")
                    return json_results
                except aiomysql.IntegrityError as e:
                    print(f"word2json表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"word2json表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f" 获取 word2json表 失败: {e}")

    # 插入 word2json表
    async def insert_word2json(self, fileId: str, word_json_dict: dict):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                word_json = json.dumps(word_json_dict, indent=4, ensure_ascii=False)
                insert_sql = """
                INSERT INTO word2json (fileId, word_json) 
                VALUES (%s, %s)
                """
                try:
                    await cursor.execute(insert_sql, (fileId, word_json))
                    await conn.commit()  # 确保提交事务
                    print("word2json表 插入气泡 成功")
                    return 'word2json表 插入气泡 成功'
                except aiomysql.IntegrityError as e:
                    print(f"word2json表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"word2json表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f"word2json表 插入气泡 失败: {e}")
                    # 可以在这里添加额外的错误处理逻辑，比如日志记录等

    # 预防最后一个批次没有出现气泡，导致进度条不为100%
    async def save_overall_progress_bar(self, fileId: str, score: int, distance: float):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                select_sql = """
                            SELECT MAX(slice_index)
                            FROM history_overall
                            WHERE fileId = %s AND (score >= %s OR distance <= %s)
                    """
                save_sql = """
                        UPDATE history_overall
                        SET bubble_num = 1
                        WHERE fileId = %s AND slice_index = %s
                    """
                try:
                    await cursor.execute(select_sql, (fileId, score, distance))
                    await conn.commit()  # 确保提交事务
                    select_comment = await cursor.fetchone()
                    select_comment_id = select_comment[0]
        
                    await cursor.execute(save_sql, (fileId, select_comment_id))
                    await conn.commit()  # 确保提交事务
                    return  '更新进度条成功'
        
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f" 更新进度条 失败: {e}")
                

    async def insert_task(self,
                          user_id: int,
                          sector: str,
                          industry: str,
                          region: str,
                          file_name: str,
                          fileId: str,
                          result: str,
                          error: str,
                          progress: str,
                          status: str,
                          total_inquiry: int,
                          current_inquiry: int,
                          token_usage: str):
        """
        :param task_id:
        :param user_id:
        :param sector:
        :param industry:
        :param region:
        :param file_name:
        :param fileId:
        :param result:
        :param error:
        :param progress:
        :param status:
        :param total_inquiry:
        :param current_inquiry:
        :return:
        """
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                # 获取当前时间
                now = datetime.now()
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
                insert_sql = """insert INTO task (user_id, sector, industry, region, file_name, file_id, created_time, updated_time, result, error, progress, status, total_inquiry, current_inquiry, token_usage) values (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"""
                try:
                    print('开始执行插入任务表的sql: %s' % insert_sql)
                    await cursor.execute(insert_sql, (user_id, sector, industry, region, file_name, fileId, update_time_str, update_time_str, result, error, progress, status, total_inquiry, current_inquiry, token_usage))
                    task_id = cursor.lastrowid
                    await conn.commit()
                    print(f" 插入任务表 成功: {task_id}")
                    return task_id
                except Exception as e:
                    await conn.rollback()
                    print(f" 插入任务表 错误: {e}")
                    raise ValueError(f" 插入任务表 错误: {e}")

    async def update_task_token_usage(self, task_id: int, token_usage: str):
        """
        更新token消耗的字段
        """
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                update_sql = """update task set token_usage = %s where task_id = %s"""
                try:
                    print('开始执行更新任务表的sql: %s' % update_sql)
                    await cursor.execute(update_sql, (token_usage, task_id))
                    await conn.commit()
                except Exception as e:
                    await conn.rollback()
                    print(f" 更新任务表 错误: {e}")
                    raise ValueError


    async def update_task_progress(self, task_id: int, total_inquiry: int, have_finished: int):
        """
        更新任务进度表
        :param task_id:
        :param progress:
        :param total_inquiry:
        :param current_inquiry:
        :return:
        """
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                # 如果已经全部完成了，那么就status变成finished
                if total_inquiry == have_finished:
                    # 任务状态变成finished
                    status = 'finished'
                else:
                    status = 'running'
                if total_inquiry == 0:
                    progress = None
                else:
                    progress = have_finished / total_inquiry
                now = datetime.now()
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
                update_sql = """update task set progress = %s, status = %s, total_inquiry = %s, current_inquiry = %s, updated_time = %s where task_id = %s"""
                try:
                    print('开始执行更新任务表的sql: %s' % update_sql)
                    await cursor.execute(update_sql, (progress, status, total_inquiry, have_finished, update_time_str, task_id))
                    await conn.commit()
                except Exception as e:
                    await conn.rollback()
                    print(f" 更新任务表 错误: {e}")
                    raise ValueError
                return '更新进度条成功'


    async def select_task(self, file_id: int):
        """
        查询任务表
        :param file_id:
        :return:
        """
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                select_sql = """select * from task where file_id = %s"""
                try:
                    print('开始执行查询任务表的sql: %s' % select_sql)
                    await cursor.execute(select_sql, (file_id,))
                    task_info = await cursor.fetchone()
                    if not task_info:
                        return None
                    task_info = {
                        'task_id': task_info[0],
                        'user_id': task_info[1],
                        'sector': task_info[2],
                        'industry': task_info[3],
                        'region': task_info[4],
                        'file_name': task_info[5],
                        'file_id': task_info[6],
                        'created_time': task_info[7],
                        'updated_time': task_info[8],
                        'progress': task_info[11],
                        'status': task_info[12],
                        'total_inquiry': task_info[13],
                        'have_finished': task_info[14],
                    }
                    return task_info
                except Exception as e:
                    print(f" 查询任务表 错误: {e}")
                    raise ValueError(f" 查询任务表 错误: {e}")

    async def update_task(self, task_id: int,
                          user_id: int,
                          sector: str,
                          industry: str,
                          region: str,
                          file_name: str,
                          fileId: str,
                          result: str,
                          error: str,
                          progress: str,
                          status: str,
                          total_inquiry: int,
                          current_inquiry: int):
        """
        更新任务表
        :param task_id:
        :param user_id:
        :param sector:
        :param industry:
        :param region:
        :param file_name:
        :param fileId:
        :param result:
        :param error:
        :param progress:
        :param status:
        :param total_inquiry:
        :param current_inquiry:
        :return:
        """
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                now = datetime.now()
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")

                update_sql = """
                UPDATE task
                SET user_id = %s, sector = %s, industry = %s, region = %s, file_name = %s, file_id = %s, updated_time = %s, result = %s, error = %s, progress = %s, status = %s, total_inquiry = %s, current_inquiry = %s
                WHERE task_id = %s
                """

                try:
                    print('开始执行更新任务表的sql: %s' % update_sql)
                    await cursor.execute(update_sql, (user_id, sector, industry, region, file_name, fileId, update_time_str, result, error, progress, status, total_inquiry, current_inquiry, task_id))
                    await conn.commit()
                    return  '更新任务表成功'
                except Exception as e:
                    await conn.rollback()
                    print(f" 更新任务表 错误: {e}")
                    raise ValueError(f" 更新任务表 错误: {e}")

    # 插入 全局表
    async def insert_history_overall(self, user_id: int, file_name: str, bubble: dict, fileId: str, status_code: int,
                             slice_index: int, score: int, distance: float):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                # 获取当前时间
                now = datetime.now()
        
                # 格式化时间为字符串
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
        
                bubble_json = json.dumps(bubble)
                insert_sql = """
                INSERT INTO history_overall (user_id, file_name, bubble, fileId, status_code, slice_index, update_time, score, distance) 
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                """
                try:
                    print('执行sql: %s' % insert_sql)
                    await cursor.execute(insert_sql, (user_id, file_name, bubble_json, fileId, status_code, slice_index, update_time_str, score, distance))
                    await conn.commit()  # 确保提交事务
                    print("全局历史表 插入气泡 成功")
                    return '全局历史表 插入气泡 成功'
                except aiomysql.IntegrityError as e:
                    print(f"全局历史表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"全局历史表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f"全局历史表 插入气泡 失败: {e}")
                    # 可以在这里添加额外的错误处理逻辑，比如日志记录等

    # 更新 全局表
    async def update_history_overall(self, fileId: str, bubble: dict, status_code: int, slice_index: int, bubble_num: float):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                        
                # 获取当前时间
                now = datetime.now()
        
                # 格式化时间为字符串
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
        
                bubble_json = json.dumps(bubble)
                update_sql = """
                UPDATE history_overall
                SET 
                    bubble = %s, 
                    status_code = %s,
                    update_time = %s,
                    bubble_num = %s
                WHERE fileId = %s AND slice_index = %s
                """
                try:
                    await cursor.execute(update_sql,
                                              (bubble_json, status_code, update_time_str, bubble_num, fileId, slice_index))
                    await conn.commit()  # 确保提交事务
                    print("全局历史表 更新气泡 成功")
                    return '全局历史表 更新气泡 成功'
                except aiomysql.IntegrityError as e:
                    print(f"全局历史表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f'全局历史表 唯一值 错误；{e}, 跳过...'
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f"全局历史表 更新气泡 失败: {e}")
                    # 可以在这里添加额外的错误处理逻辑，比如日志记录等
        
    # 获取全局 气泡字典
    async def get_history_overall_bubble(self, fileId: str, slice_index: int):
        
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                select_sql = """
                           SELECT bubble 
                           FROM history_overall
                           WHERE fileId = %s AND slice_index = %s
                           """
                try:
                    await cursor.execute(select_sql, (fileId, slice_index))
                    history_count = await cursor.fetchone()
                    if not history_count:
                        raise ValueError("未找到匹配的历史bubble数据")
                    bubble_dict = json.loads(history_count[0])
                    # problem_id_list = bubble_dict['problem_id_list']
                    print("全局历史表 获取bubble 成功")
                    return bubble_dict
                except aiomysql.IntegrityError as e:
                    print(f"全局历史表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"全局历史表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f" 获取bubble 失败: {e}")

    # 删除 全局气泡
    async def delete_history_overall_bubble(self, fileId: str, slice_index: int):

        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                delete_sql = """
                                DELETE FROM history_overall
                                WHERE fileId = %s AND slice_index = %s
                           """
                try:
                    await cursor.execute(delete_sql, (fileId, slice_index))
                    await conn.commit()  # 确保提交事务
                    # problem_id_list = bubble_dict['problem_id_list']
                    print("全局历史表 删除bubble 成功")
                except aiomysql.IntegrityError as e:
                    print(f"全局历史表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"全局历史表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f" 全局历史表 删除bubble 失败: {e}")

    ###############################################
    # 插入 文本纠错表
    async def insert_correction(self, user_id: int, bubble: dict, fileId: str, slice_index: int, progress_bar: float):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")

        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                # 获取当前时间
                now = datetime.now()

                # 格式化时间为字符串
                update_time_str = now.strftime("%Y-%m-%d %H:%M:%S")

                bubble_json = json.dumps(bubble)
                insert_sql = """
                INSERT INTO correction (user_id, bubble, fileId, slice_index, update_time, progress_bar) 
                VALUES (%s, %s, %s, %s, %s, %s)
                """
                try:
                    await cursor.execute(insert_sql, (
                        user_id, bubble_json, fileId, slice_index, update_time_str, progress_bar))
                    await conn.commit()  # 确保提交事务
                    print("文本纠错表 插入气泡 成功")
                    return '文本纠错表 插入气泡 成功'
                except aiomysql.IntegrityError as e:
                    print(f"文本纠错表 唯一值 错误；{e}, 跳过...")
                    await conn.rollback()
                    return f"文本纠错表 唯一值 错误；{e}, 跳过..."
                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f"文本纠错表 插入气泡 失败: {e}")
                    # 可以在这里添加额外的错误处理逻辑，比如日志记录等

    # 文本纠错 存入结束信息
    async def save_correction_progress_bar(self, user_id: int, fileId: str):
        if not self.pool:
            raise RuntimeError("数据库连接池未初始化")
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                select_sql = """
                            SELECT MAX(slice_index)
                            FROM correction
                            WHERE user_id = %s AND fileId = %s
                    """
                save_sql = """
                        UPDATE correction
                        SET progress_bar = 1
                        WHERE user_id = %s AND fileId = %s AND slice_index = %s
                    """
                try:
                    await cursor.execute(select_sql, (user_id, fileId))
                    await conn.commit()  # 确保提交事务
                    select_comment = await cursor.fetchone()
                    select_comment_id = select_comment[0]

                    await cursor.execute(save_sql, (user_id, fileId, select_comment_id))
                    await conn.commit()  # 确保提交事务
                    return '更新进度条成功'

                except Exception as e:
                    await conn.rollback()  # 如果有错误发生，回滚事务
                    raise ValueError(f" 更新进度条 失败: {e}")